package cn.neo.hhsharding.engine;

import cn.neo.hhsharding.config.ShardingProperties;
import cn.neo.hhsharding.demo.model.User;
import cn.neo.hhsharding.strategy.HashShardingStrategy;
import cn.neo.hhsharding.strategy.ShardingStrategy;
import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLName;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
import com.alibaba.druid.sql.ast.statement.SQLInsertStatement;
import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlSchemaStatVisitor;
import lombok.val;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;

import java.util.*;
import java.util.stream.Collectors;

/**
 * core sharding engine
 */
public class StandardShardingEngine implements ShardingEngine {

    // 库名:List<表名>
    private final MultiValueMap<String, String> actualDatabaseNames = new LinkedMultiValueMap<>();

    // 表名:List<库名>
    private final MultiValueMap<String, String> actualTableNames = new LinkedMultiValueMap<>();

    // 分库策略
    private final Map<String, ShardingStrategy> databaseStrategies = new HashMap<>();

    // 分表策略
    private final Map<String, ShardingStrategy> tableStrategies = new HashMap<>();

    public StandardShardingEngine(ShardingProperties properties) {
        properties.getTables().forEach((table, tableProperties) -> {
            tableProperties.getActualDataNodes().forEach(actualDataNode -> {
                String[] split = actualDataNode.split("\\.");
                String databaseName = split[0];
                String tableName = split[1];
                actualDatabaseNames.add(databaseName, tableName);
                actualTableNames.add(tableName, databaseName);
            });

            databaseStrategies.put(table, new HashShardingStrategy(tableProperties.getDatabaseStrategy()));
            tableStrategies.put(table, new HashShardingStrategy(tableProperties.getTableStrategy()));

        });
    }

    @Override
    public ShardingResult sharding(String sql, Object[] args) {

        SQLStatement sqlStatement = SQLUtils.parseSingleMysqlStatement(sql);
        String table;
        Map<String, Object> shardingColumnsMap = new HashMap<>();

        // insert
        if ((sqlStatement instanceof SQLInsertStatement sqlInsertStatement)) {
            // 逻辑表名
            table = sqlInsertStatement.getTableName().getSimpleName();
            // 获取列的值
            List<SQLExpr> columns = sqlInsertStatement.getColumns();
            for (int i = 0; i < columns.size(); i++) {
                SQLExpr column = columns.get(i);
                SQLIdentifierExpr columnExper = (SQLIdentifierExpr) column;
                String columnName = columnExper.getName();
                shardingColumnsMap.put(columnName, args[i]);
            }

        } else {
            // TODO select /update /delete
            // 解析语法树
            MySqlSchemaStatVisitor visitor = new MySqlSchemaStatVisitor();
            visitor.setParameters(List.of(args));
            sqlStatement.accept(visitor);

            LinkedHashSet<SQLName> sqlNames = new LinkedHashSet<>(visitor.getOriginalTables());
            if (sqlNames.size() > 1) {
                throw new RuntimeException("不支持多表查询:not support multi tables sharding:" + sqlNames);
            }

            table = sqlNames.iterator().next().getSimpleName();
            System.out.println(" ===> visitor.getOriginalTables = " + table);
            // 条件
            shardingColumnsMap = visitor.getConditions().stream()
                    .collect(Collectors.toMap(c -> c.getColumn().getName(), c -> c.getValues().get(0)));

            System.out.println(" ===> visitor.getConditions = " + shardingColumnsMap);
        }


        ShardingStrategy databaseStrategy = databaseStrategies.get(table);
        String targetDatabase = databaseStrategy.doSharding(actualDatabaseNames.get(table), table, shardingColumnsMap);

        ShardingStrategy tableStrategy = tableStrategies.get(table);
        String targetTable = tableStrategy.doSharding(actualTableNames.get(table), table, shardingColumnsMap);

        System.out.println(" ===>");
        System.out.println(" ===> target db.table = " + targetDatabase + "." + targetTable);
        System.out.println(" ===>");

        return new ShardingResult(targetDatabase, sql.replace(table, targetTable));
    }

}
